import json
import os
from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv
import DBTools as db

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPUAI_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")
client = ZhipuAI(api_key=api_key)


# 创建一个用于Function Calling的函数定义
def create_function_calling_request(messages):
    return client.chat.completions.create(
        model="glm-4-plus",
        messages=messages,
        tools=[{  
            "type": "function",
            "function": {
                "name": "ask_database",
                "description": "使用此函数来回答用户关于业务数据的问题。 \
                            应识别用户的输入并生成完整的SQL语句作为本函数的参数。",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": f"""
                            SQL查询用于提取信息以回答用户的问题。
                            SQL应该使用这个数据库模式来编写：
                            {db.db_schema}
                            查询应以纯文本形式返回，而不是JSON。
                            查询应仅包含SQLite支持的语法。
                            """,
                        }
                    },
                    "required": ["query"],
                }
            }
        }],
        tool_choice="auto"
    )


# 初始化数据库
db.initDB()

# 接收用户查询
while True:  
    user_input = input("请输入提示文字（输入'exit'退出程序）: ")  
    if user_input is None or user_input == "" :
        continue
    elif user_input.lower() == 'exit':  
        print("程序退出。")  
        break  
    else :
        messages = [{"role": "user", "content": user_input}]
        response = create_function_calling_request(messages)  
        function_call = response.choices[0].message.tool_calls[0]
        if function_call is  None:
            result = response.choices[0].message.content
        else:
            messages.append(response.choices[0].message.model_dump())
            
            function_name = function_call.function.name
            function_args = json.loads(function_call.function.arguments)
            if function_name == "ask_database":
                sql = function_args.get("query")
                print(f"\n大模型生成SQL语句：\n{sql}\n")
                result = db.ask_database(sql)
                # 将本地查询结果返回大模型
                follow_up_prompt = {
                    "role": "tool",
                    "tool_call_id": function_call.id,
                    "content": json.dumps({"result": result})
                }            
                messages.append(follow_up_prompt)        
                
                result = create_function_calling_request(messages)
                
        reply = result.choices[0].message.content.strip()
        print(f"大模型最终回复：\n{reply}\n")